from collections.abc import Callable, Iterator
from multiprocessing.pool import ThreadPool
from typing import Any, Generic, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, resume
from .state_safe import StateSafeTransformedIterator

AnyState = TypeVar('AnyState')
SrcSample = TypeVar('SrcSample')
DstSample = TypeVar('DstSample')


class ThreadedMappedIterator(DatasetIterator[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self,
        iterator: Iterator[SrcSample],
        map_func: Callable[[SrcSample], DstSample], *,
        thread_num: int,
        ordered: bool = True,
        chunk_size: int = 1,
    ):
        self._thread_pool = ThreadPool(thread_num)
        self._iterator = self._thread_pool.imap(map_func, iterator, chunk_size) \
            if ordered else self._thread_pool.imap_unordered(map_func, iterator, chunk_size)

    def __next__(self) -> DstSample:
        return next(self._iterator)

    def __del__(self):
        self._thread_pool.close()
        self._thread_pool.join()


class ThreadedMappedDataset(Dataset[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self,
        dataset: Dataset[SrcSample],
        map_func: Callable[[SrcSample], DstSample], *,
        thread_num: int = 4,
        ordered: bool = True,
    ):
        self._dataset = dataset
        self._map_func = map_func
        self._thread_num = thread_num
        self._ordered = ordered

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def map_func(self) -> Callable[[SrcSample], DstSample]:
        return self._map_func

    def __len__(self) -> int:
        return len(self._dataset)

    def __getitem__(self, index: int, /) -> DstSample:
        if not self._ordered:
            raise TypeError('ThreadedMappedDataset with unordered iteration does not support indexing!')
        return self._dataset[index]

    def __iter__(self) -> ThreadedMappedIterator[DstSample, SrcSample]:
        return ThreadedMappedIterator(
            iterator=iter(self._dataset),
            map_func=self._map_func,
            thread_num=self._thread_num,
            ordered=self._ordered)


class StateSafeThreadedMappedIterator(
    StateSafeTransformedIterator[DstSample, AnyState],
    Generic[DstSample, SrcSample, AnyState]
):
    def __init__(self,
        iterator: Iterator[SrcSample],
        map_func: Callable[[SrcSample], DstSample], *,
        thread_num: int,
        ordered: bool = True,
        chunk_size: int = 1,
    ):
        super().__init__(
            iterator,
            lambda it: ThreadedMappedIterator(
                iterator=it,
                map_func=lambda sample_and_state: (map_func(sample_and_state[0]), sample_and_state[1]),
                thread_num=thread_num,
                ordered=ordered,
                chunk_size=chunk_size))


class StateSafeThreadedMappedDataset(ThreadedMappedDataset[DstSample, SrcSample]):
    def __init__(self,
        dataset: Dataset[SrcSample],
        map_func: Callable[[SrcSample], DstSample], *,
        thread_num: int = 4,
    ):
        super().__init__(dataset, map_func, thread_num=thread_num, ordered=True)

    def __iter__(self) -> DatasetIterator[DstSample]:
        return StateSafeThreadedMappedIterator(
            iterator=iter(self._dataset),
            map_func=self.map_func,
            thread_num=self._thread_num,
            ordered=self._ordered)

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[DstSample]:
        return StateSafeThreadedMappedIterator(
            iterator=resume(self._dataset, state),
            map_func=self.map_func,
            thread_num=self._thread_num,
            ordered=self._ordered)
